import numpy as np
import pytest
from sklearn import config_context
from sklearn.linear_model import LassoCV

from skfolio.moments import ImpliedCovariance
from skfolio.prior import (
    BlackLitterman,
    EmpiricalPrior,
    FactorModel,
    LoadingMatrixRegression,
)


def test_factor_model(X, y):
    model = FactorModel()
    model.fit(X, y)
    assert model.return_distribution_
    assert model.return_distribution_.mu.shape == (20,)
    np.testing.assert_almost_equal(
        model.return_distribution_.cholesky @ model.return_distribution_.cholesky.T,
        model.return_distribution_.covariance,
        15,
    )

    model = FactorModel(
        residual_variance=False,
        loading_matrix_estimator=LoadingMatrixRegression(
            linear_regressor=LassoCV(cv=5, fit_intercept=False), n_jobs=-1
        ),
    )
    model.fit(X, y)
    assert model.return_distribution_
    np.testing.assert_almost_equal(
        model.return_distribution_.cholesky @ model.return_distribution_.cholesky.T,
        model.return_distribution_.covariance,
        15,
    )


def test_black_litterman_factor_model(X, y):
    factor_views = ["MTUM - QUAL == 0.03 ", "SIZE - USMV== 0.04", "VLUE == 0.06 "]
    n_observations = X.shape[0]
    model = FactorModel(
        factor_prior_estimator=BlackLitterman(
            views=factor_views, tau=1 / n_observations
        ),
        residual_variance=False,
    )
    model.fit(X, y)

    assert model.return_distribution_.mu.shape == (20,)
    assert model.return_distribution_.covariance.shape == (20, 20)
    np.testing.assert_almost_equal(
        model.return_distribution_.mu,
        np.array(
            [
                0.03913265,
                0.06901794,
                0.04743629,
                0.04119901,
                0.03839577,
                0.04114205,
                0.03060717,
                0.00924759,
                0.04197938,
                0.0095809,
                0.01440974,
                0.0130805,
                0.03724454,
                0.00999507,
                0.01208523,
                0.00583489,
                0.05676089,
                0.02747053,
                0.01263982,
                0.0330812,
            ]
        ),
    )

    np.testing.assert_almost_equal(
        model.return_distribution_.covariance,
        np.array(
            [
                [
                    1.93367002e-04,
                    2.54684881e-04,
                    1.73319889e-04,
                    1.63867934e-04,
                    1.42550087e-04,
                    1.51237386e-04,
                    1.46752447e-04,
                    8.70212559e-05,
                    1.56125484e-04,
                    8.49184948e-05,
                    1.03685876e-04,
                    8.89605126e-05,
                    1.94289761e-04,
                    9.26197199e-05,
                    9.24867175e-05,
                    8.15088774e-05,
                    1.52490035e-04,
                    1.31473009e-04,
                    7.28363903e-05,
                    1.23032284e-04,
                ],
                [
                    2.54684881e-04,
                    3.53587452e-04,
                    2.38238511e-04,
                    2.22055181e-04,
                    1.95549859e-04,
                    2.07424002e-04,
                    1.93869475e-04,
                    1.06874871e-04,
                    2.14106398e-04,
                    1.03836360e-04,
                    1.31145688e-04,
                    1.12862660e-04,
                    2.54689883e-04,
                    1.14177505e-04,
                    1.15748785e-04,
                    9.74320654e-05,
                    2.22472489e-04,
                    1.75212628e-04,
                    9.40612892e-05,
                    1.68723251e-04,
                ],
                [
                    1.73319889e-04,
                    2.38238511e-04,
                    2.37106293e-04,
                    1.99725893e-04,
                    1.91102869e-04,
                    2.07918604e-04,
                    1.57127674e-04,
                    9.03081333e-05,
                    2.09574635e-04,
                    1.02340858e-04,
                    1.05446895e-04,
                    9.48912097e-05,
                    1.68644484e-04,
                    1.00487930e-04,
                    1.07848466e-04,
                    8.85425748e-05,
                    2.33840855e-04,
                    1.48885222e-04,
                    7.94266713e-05,
                    1.65152065e-04,
                ],
                [
                    1.63867934e-04,
                    2.22055181e-04,
                    1.99725893e-04,
                    1.75461202e-04,
                    1.62136366e-04,
                    1.75723711e-04,
                    1.45316891e-04,
                    8.51154732e-05,
                    1.77661262e-04,
                    9.47939163e-05,
                    9.97123601e-05,
                    8.90981301e-05,
                    1.62788169e-04,
                    9.58201946e-05,
                    9.65698347e-05,
                    8.38001993e-05,
                    1.88697139e-04,
                    1.35265930e-04,
                    7.41116683e-05,
                    1.40003222e-04,
                ],
                [
                    1.42550087e-04,
                    1.95549859e-04,
                    1.91102869e-04,
                    1.62136366e-04,
                    1.54262333e-04,
                    1.67646867e-04,
                    1.29056105e-04,
                    7.48067867e-05,
                    1.69152283e-04,
                    8.43379427e-05,
                    8.74721755e-05,
                    7.85285826e-05,
                    1.39548760e-04,
                    8.36184674e-05,
                    8.84675362e-05,
                    7.35620323e-05,
                    1.86421863e-04,
                    1.22451848e-04,
                    6.55536362e-05,
                    1.33297852e-04,
                ],
                [
                    1.51237386e-04,
                    2.07424002e-04,
                    2.07918604e-04,
                    1.75723711e-04,
                    1.67646867e-04,
                    1.82611772e-04,
                    1.38755136e-04,
                    8.00373229e-05,
                    1.83825440e-04,
                    9.16001420e-05,
                    9.32188124e-05,
                    8.41445402e-05,
                    1.47503759e-04,
                    8.97225414e-05,
                    9.50382757e-05,
                    7.89551361e-05,
                    2.04261264e-04,
                    1.31157444e-04,
                    7.04452536e-05,
                    1.44860808e-04,
                ],
                [
                    1.46752447e-04,
                    1.93869475e-04,
                    1.57127674e-04,
                    1.45316891e-04,
                    1.29056105e-04,
                    1.38755136e-04,
                    1.30054400e-04,
                    8.15744627e-05,
                    1.41304393e-04,
                    8.84139812e-05,
                    9.56669932e-05,
                    8.43448378e-05,
                    1.50993014e-04,
                    9.40810770e-05,
                    8.76330985e-05,
                    8.22017170e-05,
                    1.33455354e-04,
                    1.22627131e-04,
                    6.87833693e-05,
                    1.11352751e-04,
                ],
                [
                    8.70212559e-05,
                    1.06874871e-04,
                    9.03081333e-05,
                    8.51154732e-05,
                    7.48067867e-05,
                    8.00373229e-05,
                    8.15744627e-05,
                    5.84367010e-05,
                    8.19463902e-05,
                    6.28559931e-05,
                    6.71329448e-05,
                    5.86002506e-05,
                    9.30778794e-05,
                    6.89760262e-05,
                    6.16286240e-05,
                    6.16329454e-05,
                    6.36869852e-05,
                    8.09541082e-05,
                    4.62469233e-05,
                    6.45765912e-05,
                ],
                [
                    1.56125484e-04,
                    2.14106398e-04,
                    2.09574635e-04,
                    1.77661262e-04,
                    1.69152283e-04,
                    1.83825440e-04,
                    1.41304393e-04,
                    8.19463902e-05,
                    1.85484550e-04,
                    9.23429765e-05,
                    9.58021475e-05,
                    8.59922127e-05,
                    1.52781380e-04,
                    9.15378305e-05,
                    9.70054484e-05,
                    8.05722501e-05,
                    2.04506529e-04,
                    1.34148068e-04,
                    7.17717303e-05,
                    1.46168244e-04,
                ],
                [
                    8.49184948e-05,
                    1.03836360e-04,
                    1.02340858e-04,
                    9.47939163e-05,
                    8.43379427e-05,
                    9.16001420e-05,
                    8.84139812e-05,
                    6.28559931e-05,
                    9.23429765e-05,
                    7.25017552e-05,
                    7.12438214e-05,
                    6.36919321e-05,
                    9.12395170e-05,
                    7.64539778e-05,
                    6.66689077e-05,
                    6.79358212e-05,
                    7.55967219e-05,
                    8.72297938e-05,
                    5.07844289e-05,
                    7.27694610e-05,
                ],
                [
                    1.03685876e-04,
                    1.31145688e-04,
                    1.05446895e-04,
                    9.97123601e-05,
                    8.74721755e-05,
                    9.32188124e-05,
                    9.56669932e-05,
                    6.71329448e-05,
                    9.58021475e-05,
                    7.12438214e-05,
                    7.80984846e-05,
                    6.79277733e-05,
                    1.10935014e-04,
                    7.92761341e-05,
                    7.08158342e-05,
                    7.02101381e-05,
                    7.55156887e-05,
                    9.53832054e-05,
                    5.37817160e-05,
                    7.54954074e-05,
                ],
                [
                    8.89605126e-05,
                    1.12862660e-04,
                    9.48912097e-05,
                    8.90981301e-05,
                    7.85285826e-05,
                    8.41445402e-05,
                    8.43448378e-05,
                    5.86002506e-05,
                    8.59922127e-05,
                    6.36919321e-05,
                    6.79277733e-05,
                    5.95691953e-05,
                    9.50206634e-05,
                    6.97509751e-05,
                    6.19955548e-05,
                    6.16170919e-05,
                    6.99001499e-05,
                    8.36142120e-05,
                    4.74070832e-05,
                    6.77648394e-05,
                ],
                [
                    1.94289761e-04,
                    2.54689883e-04,
                    1.68644484e-04,
                    1.62788169e-04,
                    1.39548760e-04,
                    1.47503759e-04,
                    1.50993014e-04,
                    9.30778794e-05,
                    1.52781380e-04,
                    9.12395170e-05,
                    1.10935014e-04,
                    9.50206634e-05,
                    1.98686422e-04,
                    1.01822834e-04,
                    9.71374916e-05,
                    8.93088013e-05,
                    1.39287321e-04,
                    1.38117373e-04,
                    7.70447741e-05,
                    1.20397015e-04,
                ],
                [
                    9.26197199e-05,
                    1.14177505e-04,
                    1.00487930e-04,
                    9.58201946e-05,
                    8.36184674e-05,
                    8.97225414e-05,
                    9.40810770e-05,
                    6.89760262e-05,
                    9.15378305e-05,
                    7.64539778e-05,
                    7.92761341e-05,
                    6.97509751e-05,
                    1.01822834e-04,
                    8.42349788e-05,
                    7.19750896e-05,
                    7.45557719e-05,
                    6.62417846e-05,
                    9.54774947e-05,
                    5.49516600e-05,
                    7.21349781e-05,
                ],
                [
                    9.24867175e-05,
                    1.15748785e-04,
                    1.07848466e-04,
                    9.65698347e-05,
                    8.84675362e-05,
                    9.50382757e-05,
                    8.76330985e-05,
                    6.16286240e-05,
                    9.70054484e-05,
                    6.66689077e-05,
                    7.08158342e-05,
                    6.19955548e-05,
                    9.71374916e-05,
                    7.19750896e-05,
                    6.77456125e-05,
                    6.46261663e-05,
                    8.39668332e-05,
                    8.83940109e-05,
                    4.91545176e-05,
                    7.64436502e-05,
                ],
                [
                    8.15088774e-05,
                    9.74320654e-05,
                    8.85425748e-05,
                    8.38001993e-05,
                    7.35620323e-05,
                    7.89551361e-05,
                    8.22017170e-05,
                    6.16329454e-05,
                    8.05722501e-05,
                    6.79358212e-05,
                    7.02101381e-05,
                    6.16170919e-05,
                    8.93088013e-05,
                    7.45557719e-05,
                    6.46261663e-05,
                    6.67076791e-05,
                    5.73620263e-05,
                    8.35426618e-05,
                    4.82593764e-05,
                    6.34937212e-05,
                ],
                [
                    1.52490035e-04,
                    2.22472489e-04,
                    2.33840855e-04,
                    1.88697139e-04,
                    1.86421863e-04,
                    2.04261264e-04,
                    1.33455354e-04,
                    6.36869852e-05,
                    2.04506529e-04,
                    7.55967219e-05,
                    7.55156887e-05,
                    6.99001499e-05,
                    1.39287321e-04,
                    6.62417846e-05,
                    8.39668332e-05,
                    5.73620263e-05,
                    2.57940779e-04,
                    1.20790864e-04,
                    6.16534433e-05,
                    1.61158222e-04,
                ],
                [
                    1.31473009e-04,
                    1.75212628e-04,
                    1.48885222e-04,
                    1.35265930e-04,
                    1.22451848e-04,
                    1.31157444e-04,
                    1.22627131e-04,
                    8.09541082e-05,
                    1.34148068e-04,
                    8.72297938e-05,
                    9.53832054e-05,
                    8.36142120e-05,
                    1.38117373e-04,
                    9.54774947e-05,
                    8.83940109e-05,
                    8.35426618e-05,
                    1.20790864e-04,
                    1.22665914e-04,
                    6.72621284e-05,
                    1.05713320e-04,
                ],
                [
                    7.28363903e-05,
                    9.40612892e-05,
                    7.94266713e-05,
                    7.41116683e-05,
                    6.55536362e-05,
                    7.04452536e-05,
                    6.87833693e-05,
                    4.62469233e-05,
                    7.17717303e-05,
                    5.07844289e-05,
                    5.37817160e-05,
                    4.74070832e-05,
                    7.70447741e-05,
                    5.49516600e-05,
                    4.91545176e-05,
                    4.82593764e-05,
                    6.16534433e-05,
                    6.72621284e-05,
                    3.80674507e-05,
                    5.65586071e-05,
                ],
                [
                    1.23032284e-04,
                    1.68723251e-04,
                    1.65152065e-04,
                    1.40003222e-04,
                    1.33297852e-04,
                    1.44860808e-04,
                    1.11352751e-04,
                    6.45765912e-05,
                    1.46168244e-04,
                    7.27694610e-05,
                    7.54954074e-05,
                    6.77648394e-05,
                    1.20397015e-04,
                    7.21349781e-05,
                    7.64436502e-05,
                    6.34937212e-05,
                    1.61158222e-04,
                    1.05713320e-04,
                    5.65586071e-05,
                    1.15185634e-04,
                ],
            ]
        ),
    )


def test_metadata_routing_error(X, y, implied_vol):
    with config_context(enable_metadata_routing=True):
        model = FactorModel(
            factor_prior_estimator=EmpiricalPrior(
                covariance_estimator=ImpliedCovariance().set_fit_request(
                    implied_vol=True
                )
            )
        )

        with pytest.raises(
            ValueError, match="The following assets are missing from `implied_vol`"
        ):
            model.fit(X, y, implied_vol=implied_vol)


def test_metadata_routing(X, implied_vol):
    with config_context(enable_metadata_routing=True):
        model = FactorModel(
            factor_prior_estimator=EmpiricalPrior(
                covariance_estimator=ImpliedCovariance().set_fit_request(
                    implied_vol=True
                )
            )
        )

        with pytest.raises(ValueError):
            model.fit(X, X)

        model.fit(X, X, implied_vol=implied_vol)

    # noinspection PyUnresolvedReferences
    assert model.factor_prior_estimator_.covariance_estimator_.r2_scores_.shape == (20,)
